import os
import numpy as np
import cv2
import random


def IoU(coor1, coor2):
    """
    Arguments:
      coor1 (tuple): xmin, ymin, xmax, ymax
      coor2 (tuple): xmin, ymin, xmax, ymax
    """
    x1, y1, x2, y2 = coor1
    x3, y3, x4, y4 = coor2
    x_overlap = max(0, min(x2, x4) - max(x1, x3))
    y_overlap = max(0, min(y2, y4) - max(y1, y3))
    intersection = x_overlap * y_overlap
    return intersection > 0


def read_points(points_path):
    with open(os.path.join(points_path), "r") as f:
        f.readline()  # 去除标题行
        for i, line in enumerate(f):
            x, y, idx = line.split(",")
            if i == 0:
                points = [(int(x), int(y))]
                idx_last = idx
            elif idx == idx_last:
                points.append((int(x), int(y)))
            else:
                yield points
                points = [(int(x), int(y))]
                idx_last = idx


def extract_subfig(collect_dir):
    """
    不需要对陨石坑作正射校正，直接按检测框拿出来就可
    """
    for sub_dir in os.listdir(collect_dir):
        # 子文件夹一级
        points_dir = os.path.join(collect_dir, sub_dir, "points_id")
        img_dir = os.path.join(collect_dir, sub_dir, "images")
        for file in os.listdir(img_dir):
            # 图像文件夹一级
            img = cv2.imread(os.path.join(img_dir, file))
            windows = []
            points_file = os.path.join(points_dir, file.replace(".png", ".txt"))
            if not os.path.exists(points_file):
                continue
            for points in read_points(points_file):
                points = np.array(points)
                x_min, y_min = points.min(axis=0)
                x_max, y_max = points.max(axis=0)
                windows.append((x_min, y_min, x_max, y_max))
            for xmin, ymin, xmax, ymax in windows:
                shape = min(xmax - xmin, ymax - ymin)
                margin = shape // 3
                if (
                    xmin - margin > 0
                    and ymin - margin > 0
                    and xmax + margin < img.shape[1]
                    and ymax + margin < img.shape[0]
                ):
                    crater = img[
                        ymin - margin : ymax + margin, xmin - margin : xmax + margin, :
                    ]
                    if crater.size == 0:
                        continue
                    i = 0
                    error = False
                    while True:
                        i += 1
                        if i > 100:
                            error = True
                            break
                        x = random.randint(shape, img.shape[0] - shape)
                        y = random.randint(shape, img.shape[1] - shape)
                        if any(
                            [
                                IoU(
                                    (
                                        x - shape // 2,
                                        y - shape // 2,
                                        x + shape // 2,
                                        y + shape // 2,
                                    ),
                                    xyxy,
                                )
                                for xyxy in windows
                            ]
                        ):
                            continue
                        if (
                            x - shape // 2 > 0
                            and y - shape // 2 > 0
                            and x + shape // 2 < img.shape[0]
                            and y + shape // 2 < img.shape[1]
                        ):
                            non_crater = img[
                                x - shape // 2 : x + shape // 2,
                                y - shape // 2 : y + shape // 2,
                                :,
                            ]
                            if non_crater.size == 0:
                                continue
                            else:
                                break
                    if not error:
                        yield crater, non_crater


if __name__ == "__main__":
    collect_dir = "/disk527/sdb1/a804_cbf/datasets/collect_data"
    output_dir = "/disk527/sdb1/a804_cbf/datasets/gazebo/craters"
    if not os.path.exists(f"{output_dir}/craters/32x32"):
        os.makedirs(f"{output_dir}/craters/32x32")
        os.makedirs(f"{output_dir}/non_craters/32x32")
    train = open(os.path.join(output_dir, "train.txt"), "w")
    val = open(os.path.join(output_dir, "val.txt"), "w")
    for it, (crater, non_crater) in enumerate(extract_subfig(collect_dir)):
        crater = cv2.resize(crater, (32, 32))
        non_crater = cv2.resize(non_crater, (32, 32))
        cv2.imwrite(f"{output_dir}/craters/32x32/{it}.png", crater)
        cv2.imwrite(f"{output_dir}/non_craters/32x32/{it}.png", non_crater)
        if random.random() < 0.1:
            print(f"craters/32x32/{it}.png", file=val, flush=True)
        else:
            print(f"craters/32x32/{it}.png", file=train, flush=True)
        if random.random() < 0.1:
            print(f"non_craters/32x32/{it}.png", file=val, flush=True)
        else:
            print(f"non_craters/32x32/{it}.png", file=train, flush=True)
        if it > 20000:
            break
    train.close()
    val.close()
